import os
import xml.etree.ElementTree as ET
from xml.etree.ElementTree import Element

def __indent(elem, level=0):
    i = "\n" + level * "\t"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "\t"
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            __indent(elem, level + 1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i
def write_xml(path,results,save_xml_folder='./',width=640,height=480):
    print(path)
    image_basename = os.path.basename(path)
    text = image_basename.split('.')[0]
    root = ET.Element('annotation')
    tree = ET.ElementTree(root)
    folder = Element("folder")
    folder.text="R"
    conpoment = Element("conpoment")
    if len(results['cls'])>0:
        result1 = results['cls'][0]
        result2 = results['cls'][1]
    else:
        result1 = 'None'
        result2 ='None'
    conpoment.text =result1 #result1
    ng_type = Element("ng_type")
    ng_type.text = result2 #result2
    filename = ET.Element('filename')
    filename.text=image_basename
    pathnode = ET.Element('path')
    pathnode.text=str(path)

    source = ET.Element('source')
    database = ET.Element('database')
    database.text="Unknow"
    source.append(database)

    size = ET.Element('size')
    widthnode = ET.Element('width')
    widthnode.text = str(width)
    heightnode = ET.Element('height')
    heightnode.text= str(height)
    depthnode = ET.Element('depth')
    depthnode.text = str(3)
    size.append(widthnode)
    size.append(heightnode)
    size.append(depthnode)

    segmented = ET.Element('segmented')
    segmented.text = '0'

    root.append(folder)
    root.append(filename)
    root.append(conpoment)
    root.append(ng_type)
    root.append(pathnode)
    root.append(source)
    root.append(size)
    root.append(segmented)

    res_yolo = results['res']
    for label,conf,box in res_yolo:
        y1, x1, y2, x2 = box
        #x:141,y1:340,x2171,y2:381
        object = ET.Element('object')
        name = ET.Element('name')
        name.text=label
        pose = ET.Element('pose')
        pose.text = 'Unspecified'
        truncated = ET.Element('truncated')
        truncated.text = '0'
        difficult = ET.Element('difficult')
        difficult.text = '0'

        bndbox = ET.Element('bndbox')
        xmin = ET.Element('xmin')
        xmin.text = str(float(x1))
        ymin = ET.Element('ymin')
        ymin.text = str(float(y1))
        xmax = ET.Element('xmax')
        xmax.text = str(float(x2))
        ymax = ET.Element('ymax')
        ymax.text = str(float(y2))
        bndbox.append(xmin)
        bndbox.append(ymin)
        bndbox.append(xmax)
        bndbox.append(ymax)

        object.append(name)
        object.append(pose)
        object.append(truncated)
        object.append(difficult)
        object.append(bndbox)
        root.append(object)

    __indent(root)

    tree.write(os.path.join(save_xml_folder,text+'.xml'),encoding='utf-8',xml_declaration=True)
